import os
import subprocess

class VideoMasker:
    """
    给视频四边添加矩形或进行裁剪
    """

    def __init__(self, input_folder=None, output_folder=r"C:\Users\33720\Desktop",
                 top=0, bottom=0, left=0, right=0, mode="mask"):
        """
        :param input_folder: 输入视频文件夹
        :param output_folder: 输出文件夹
        :param top, bottom, left, right: 四个方向的像素值
        :param mode: 'mask' -> 添加矩形, 'crop' -> 裁剪视频
        """
        self.input_folder = input_folder
        self.output_folder = output_folder
        self.top = top
        self.bottom = bottom
        self.left = left
        self.right = right
        self.mode = mode.lower()
        os.makedirs(self.output_folder, exist_ok=True)
        self.supported_exts = ('.mp4', '.avi', '.mov', '.mkv', '.ts')

    # ==================== Setter 方法 ====================
    def set_input_folder(self, path):
        self.input_folder = path

    def set_output_folder(self, path):
        self.output_folder = path
        os.makedirs(self.output_folder, exist_ok=True)

    def set_top(self, pixels):
        self.top = pixels

    def set_bottom(self, pixels):
        self.bottom = pixels

    def set_left(self, pixels):
        self.left = pixels

    def set_right(self, pixels):
        self.right = pixels

    def set_mode(self, mode):
        if mode.lower() in ["mask", "crop"]:
            self.mode = mode.lower()
        else:
            print(f"❌ 模式错误: {mode}, 保持原模式 {self.mode}")

    # ==================== 核心处理函数 ====================
    def process_video(self, input_path, output_path):
        """
        处理单个视频
        """
        if self.mode == "mask":
            draw_filters = []
            if self.top > 0:
                draw_filters.append(f"drawbox=x=0:y=0:w=iw:h={self.top}:color=white@1:t=fill")
            if self.bottom > 0:
                draw_filters.append(f"drawbox=x=0:y=ih-{self.bottom}:w=iw:h={self.bottom}:color=white@1:t=fill")
            if self.left > 0:
                draw_filters.append(f"drawbox=x=0:y=0:w={self.left}:h=ih:color=white@1:t=fill")
            if self.right > 0:
                draw_filters.append(f"drawbox=x=iw-{self.right}:y=0:w={self.right}:h=ih:color=white@1:t=fill")
            vf = ",".join(draw_filters)

        elif self.mode == "crop":
            vf = f"crop=iw-{self.left+self.right}:ih-{self.top+self.bottom}:{self.left}:{self.top}"
        else:
            print(f"❌ 模式错误: {self.mode}")
            return

        cmd = [
            "ffmpeg", "-y", "-i", input_path,
            "-vf", vf,
            "-c:a", "copy",
            output_path
        ]

        print(f"正在处理: {input_path}")
        subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        return output_path

    def process_all(self):
        """
        遍历文件夹及子文件夹处理所有视频
        """
        result = []
        if not self.input_folder or not os.path.exists(self.input_folder):
            print(f"❌ 输入文件夹不存在: {self.input_folder}")
            return result

        for root, _, files in os.walk(self.input_folder):
            rel_path = os.path.relpath(root, self.input_folder)
            current_output_folder = os.path.join(self.output_folder, rel_path)
            os.makedirs(current_output_folder, exist_ok=True)

            for filename in files:
                if filename.lower().endswith(self.supported_exts):
                    input_path = os.path.join(root, filename)
                    name, ext = os.path.splitext(filename)
                    output_path = os.path.join(current_output_folder, f"{name}_{self.mode}{ext}")
                    self.process_video(input_path, output_path)
                    result.append(output_path)
        return result
